from typing import Any

from opik_optimizer import (
    ChatPrompt,
    ParameterOptimizer,
    ParameterSearchSpace,
)
from opik_optimizer.datasets import hotpot

from opik.evaluation.metrics import LevenshteinRatio
from opik.evaluation.metrics.score_result import ScoreResult

from adk_agent import ADKAgent

dataset = hotpot(count=300)


def levenshtein_ratio(dataset_item: dict[str, Any], llm_output: str) -> ScoreResult:
    metric = LevenshteinRatio()
    return metric.score(reference=dataset_item["answer"], output=llm_output)


system_prompt = """
You are a helpful assistant. Use the `search_wikipedia` tool to find factual information when appropriate.
The user will provide a question string like "Who is Barack Obama?".
1. Extract the item to look up
2. Use the `search_wikipedia` tool to find details
3. Respond clearly to the user, stating the answer found by the tool.
"""

prompt = ChatPrompt(
    system=system_prompt,
    user="{question}",
)

# Optimize it:
optimizer = ParameterOptimizer(
    model="openai/gpt-4o-mini",
    default_n_trials=20,
    n_threads=1,
    seed=42,
)

parameter_space = ParameterSearchSpace.model_validate(
    {
        "temperature": {"type": "float", "min": 0.0, "max": 1.0},
        "top_p": {"type": "float", "min": 0.3, "max": 1.0},
        "frequency_penalty": {
            "type": "float",
            "min": -1.0,
            "max": 1.0,
        },
    }
)

optimization_result = optimizer.optimize_parameter(
    prompt=prompt,
    agent_class=ADKAgent,
    dataset=dataset,
    metric=levenshtein_ratio,
    parameter_space=parameter_space,
    max_trials=2,
    n_samples=10,
)

optimization_result.display()
